import time

from backend.v3_6 import *
import torch
from torch import nn
from torch.nn import functional as F
from timm.models import hrnet
from collections import OrderedDict
BatchNorm2d = nn.BatchNorm2d

def upsample(x, scale_factor=2, mode="bilinear", align_corners=False):
    """Upsample input tensor by a factor of 2
    """
    return F.interpolate(x, scale_factor=scale_factor, mode=mode, align_corners=align_corners)

def up_cat(x_low, x):
    x_low = F.interpolate(x_low, scale_factor=2,mode='bilinear')
    return torch.cat([x_low, x],dim=1)

def up_add(x_low, x):
    x_low = F.interpolate(x_low, scale_factor=2,mode='bilinear')
    return x_low+x


'''

def up_cat(x_low, x):
    h, w = x.size(2), x.size(3)
    x_low = F.interpolate(x_low, size=(h,w),mode='bilinear')
    return torch.cat([x_low, x],dim=1)
'''

class ConvNormAct(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=(0, 0), dilation=(1, 1), bias=False):
        super(ConvNormAct, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, padding=padding, stride=stride,
                      bias=bias,dilation=dilation),
            BatchNorm2d(out_channels, momentum=0.1),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        x = self.conv(x)
        return x


class SegUNet(nn.Module):
    def __init__(self,num_class,   load_pretrained=False, is_edge=False):
        super(SegUNet, self).__init__()
        self.num_class = num_class
        self.load_pretrained = load_pretrained

        self.base_channel = 16
        self.is_edge = is_edge
        self.stem = nn.Sequential(
            ConvNormAct(3, self.base_channel, kernel_size=7, stride=2, padding=3),
            ConvNormAct(self.base_channel, self.base_channel*2, kernel_size=3, stride=2, padding=1)
        )

        self.conv1_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
        )
        self.conv1_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
        )
        self.conv1_3 = nn.Sequential(
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
        )
        self.conv1_4 = nn.Sequential(
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 2, self.base_channel * 2, kernel_size=3, stride=1, padding=1),
        )

        self.conv2_1 = nn.Sequential(

            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
        )
        self.conv2_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
        )
        self.conv2_3 = nn.Sequential(
            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 4, self.base_channel * 4, kernel_size=3, stride=1, padding=1),
        )

        self.conv3_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 8, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 8, self.base_channel * 8, kernel_size=3, stride=1, padding=1),
        )
        self.conv3_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 8, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 8, self.base_channel * 8, kernel_size=3, stride=1, padding=1),
        )

        self.conv4_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 16, self.base_channel * 16, kernel_size=3, stride=1, padding=1),
            ConvNormAct(self.base_channel * 16, self.base_channel * 16, kernel_size=3, stride=1, padding=1),
        )

        self.down_1_2_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel*2, self.base_channel*4, kernel_size=1, stride=1)
        )
        self.down_1_2_2 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel*2, self.base_channel * 4, kernel_size=1, stride=1)
        )
        self.down_1_2_3 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel*2, self.base_channel * 4, kernel_size=1, stride=1)
        )

        self.down_1_3_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 4, self.base_channel * 8, kernel_size=1, stride=1)
        )
        self.down_1_3_2 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 4, self.base_channel * 8, kernel_size=1, stride=1)
        )

        self.down_1_4_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 8, self.base_channel * 16, kernel_size=1, stride=1)
        )

        self.down_2_3_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 4, self.base_channel * 8, kernel_size=1, stride=1)
        )
        self.down_2_3_2 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 4, self.base_channel * 8, kernel_size=1, stride=1)
        )

        self.down_2_4_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 8, self.base_channel * 16, kernel_size=1, stride=1)
        )

        self.down_3_4_1 = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),
            ConvNormAct(self.base_channel * 8, self.base_channel * 16, kernel_size=1, stride=1)
        )

        self.up_2_1_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 4, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_2_1_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 4, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_2_1_3 = nn.Sequential(
            ConvNormAct(self.base_channel * 4, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )

        self.up_3_1_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_3_1_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_4_1_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 16, self.base_channel * 2, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )

        self.up_3_2_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 4, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )
        self.up_3_2_2 = nn.Sequential(
            ConvNormAct(self.base_channel * 8, self.base_channel * 4, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),

        )
        self.up_4_2_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 16, self.base_channel * 4, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )


        self.up_4_3_1 = nn.Sequential(
            ConvNormAct(self.base_channel * 16, self.base_channel * 8, kernel_size=1, stride=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
        )


    def forward(self, x):
        x1 = self.stem(x)
        x1 = self.conv1_1(x1)

        x2 = self.down_1_2_1(x1)

        x2 = self.conv2_1(x2)
        x1 = self.conv1_2(x1)

        x1_2 = self.down_1_2_2(x1)

        x1_3 = self.down_1_3_1(x1)
        x2_3 = self.down_2_3_1(x2)

        x_2_1 = self.up_2_1_1(x2)

        x1 = self.conv1_3(x1+x_2_1)
        x2 = self.conv2_2(x2+x1_2)
        x3 = self.conv3_1(x1_3+x2_3)


        x1_2 = self.down_1_2_3(x1)
        x1_3 = self.down_1_3_2(x1)
        x1_4 = self.down_1_4_1(x1)


        x2_3 = self.down_2_3_2(x2)
        x2_4 = self.down_2_4_1(x2)

        x3_4 = self.down_3_4_1(x3)

        x2_1 = self.up_2_1_2(x2)
        x3_1 = self.up_3_1_1(x3)


















if __name__ == '__main__':
    import torchstat

    model = SegUNet(num_class=10, load_pretrained=False)
    model = model.cuda()
    model.eval()
    xs = torch.randn(1, 3, 544, 960).cuda()
    print(xs.size())
    m = model(xs)
    print(m.size())
